from typing import Optional
from sqlmodel import SQLModel, Field, create_engine, Session
import pandas as pd


# 1. 定义数据模型（带中文描述）
class EnvironmentalMonitoring(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True, description="主键ID")
    year: Optional[str] = Field(default=None, max_length=10, description="年份")
    base_station_id: Optional[str] = Field(
        default=None, max_length=50, description="通信基站编号"
    )
    base_station_name: Optional[str] = Field(
        default=None, max_length=100, description="通信基站名"
    )
    project_name: Optional[str] = Field(
        default=None, max_length=100, description="所属项目"
    )
    network_type: Optional[str] = Field(
        default=None, max_length=50, description="网络制式"
    )
    frequency_range: Optional[str] = Field(
        default=None, max_length=50, description="频率范围(MHz)"
    )
    address: Optional[str] = Field(
        default=None, max_length=255, description="通信基站地址"
    )
    sector: Optional[str] = Field(default=None, max_length=50, description="扇区")
    tower_type: Optional[str] = Field(
        default=None, max_length=50, description="杆塔类型"
    )
    height: Optional[str] = Field(default=None, max_length=20, description="高度")
    vertical_distance: Optional[str] = Field(
        default=None, max_length=20, description="垂直距离"
    )
    longitude: Optional[float] = Field(default=None, description="经度")
    latitude: Optional[float] = Field(default=None, description="纬度")
    district: Optional[str] = Field(default=None, max_length=50, description="所属区县")
    construction_unit: Optional[str] = Field(
        default=None, max_length=50, description="建设单位"
    )
    remarks: Optional[str] = Field(default=None, max_length=100, description="备注")
    town: Optional[str] = Field(default=None, max_length=50, description="镇区")
    area_division: Optional[str] = Field(
        default=None, max_length=50, description="区分"
    )
    company_area_division: Optional[str] = Field(
        default=None, max_length=50, description="公司区域划分"
    )


# 2. 数据库连接配置
DATABASE_URL = "mysql+pymysql://username:password@localhost:3306/database_name"
engine = create_engine(DATABASE_URL)


# 3. 创建表
def create_db_and_tables():
    SQLModel.metadata.create_all(engine)


# 4. 从Excel导入数据
def import_from_excel(file_path: str):
    # 读取Excel文件
    df = pd.read_excel(file_path, sheet_name="Sheet1")

    # 重命名列名（中文转英文）
    column_mapping = {
        "年份": "year",
        "通信基站编号": "base_station_id",
        "通信基站名": "base_station_name",
        "所属项目": "project_name",
        "网络制式": "network_type",
        "频率范围(MHz)": "frequency_range",
        "通信基站地址": "address",
        "扇区": "sector",
        "杆塔类型": "tower_type",
        "高度": "height",
        "垂直距离": "vertical_distance",
        "经度": "longitude",
        "纬度": "latitude",
        "所属区县": "district",
        "建设单位": "construction_unit",
        "备注": "remarks",
        "镇区": "town",
        "区分": "area_division",
        "公司区域划分": "company_area_division",
    }
    df.rename(columns=column_mapping, inplace=True)

    # 处理数据
    df = df.where(pd.notnull(df), None)  # 将NaN替换为None

    # 转换为字典列表
    records = df.to_dict(orient="records")

    # 批量插入数据
    with Session(engine) as session:
        for record in records:
            # 处理经度纬度，确保是数字类型
            if "longitude" in record and pd.notna(record["longitude"]):
                record["longitude"] = float(record["longitude"])
            else:
                record["longitude"] = None

            if "latitude" in record and pd.notna(record["latitude"]):
                record["latitude"] = float(record["latitude"])
            else:
                record["latitude"] = None

            monitoring = EnvironmentalMonitoring(**record)
            session.add(monitoring)

        session.commit()


# 主函数
def main():
    # 创建数据库表
    create_db_and_tables()

    # 导入Excel数据
    excel_file = "24年环评监测清单.xlsx"
    try:
        import_from_excel(excel_file)
        print(f"成功从 {excel_file} 导入数据到数据库")
    except Exception as e:
        print(f"导入数据时出错: {e}")


if __name__ == "__main__":
    main()
